#====================================
# zefMoves - Animation copying tools
# ToolKit - zefToolKit
# version: .1
# author:  Matt Murray
#====================================

import maya.cmds as cmds
#from functools import partial
import xml.dom.minidom as xd
import os


class ZefMotion():
    
    def __init__(self):
        #this is the contructor
        
        #GLOBALS
        self.rootDir = os.getcwd()
        self.winWidth = 215
        self.winHeight = 200
        self.UIElements = {}
        self.ver = '0.1'
        
        self.sourceObjects = []
        self.destObjects = []
        self.useRange = False
        self.useType = False
        self.type = ''
        self.findSelected = True
        self.findAll = False
        self.saveByType = 'none'
        
        #buildUI
        self.buildUI()
        
    def buildUI(self):
        """Create all the main UI elements for the tool."""
        #create the main window
        self.UIElements["mainWindow"] = cmds.window(title="zefMoves - Animation Tool v"+self.ver, wh=[self.winWidth, self.winHeight], rtf=True)
        self.UIElements["mainLayout"] = cmds.rowColumnLayout(cw=[(1,215)])
        self.UIElements["mainTabs"] = cmds.tabLayout(innerMarginWidth=5, innerMarginHeight=5)
        
        #create tabs
        
        #=====Save Tab
        self.UIElements["saveTab"] = cmds.rowColumnLayout(nr=11, ral=[(1,'left'),(2,'left')], rs=[(1,15), (2,15), (3,15), (4,15)], w=100)
        self.UIElements["saveWhatCB"] = cmds.checkBoxGrp(ncb=2, cw=[(1, 175), (2,50)], vr=True, labelArray2=['Save All Animation in Scene','Save Animation for Selected'], v2=True, cc1=self.setSaveAll, cc2=self.setSaveSel)
        self.UIElements["rangeCB"] = cmds.checkBox(l="Use range(disable to eval all)", cc=self.setRange)
        self.UIElements["saveStartFrame"] = cmds.textFieldGrp(l='Start Frame: ', tx='1', en=False, cal=[(1,'left')], cw=[(1,75), (2,125)])
        self.UIElements["saveEndFrame"] = cmds.textFieldGrp(l='End Frame: ', tx='100', en=False, cal=[(1,'left')], cw=[(1,75), (2,125)])
        cmds.separator()
        self.UIElements["saveTypeCB"] = cmds.checkBox(l='Save by Type:', cc=self.setType)
        self.UIElements["saveTypeMenu"] = cmds.optionMenu(l='Type: ', en=False, w=205, cc=self.storeType)
        cmds.menuItem(l='Bones')
        cmds.menuItem(l='Curves')
        cmds.menuItem(l='Polys')
        cmds.menuItem(l='Nurbs')
        cmds.separator()
        self.UIElements["saveAnimationBtn"] = cmds.button(l='SAVE ANIMATION', bgc=(.3, .7, .9), c=self.saveAnimData)
        self.UIElements["loadAnimationBtn"] = cmds.button(l='LOAD ANIMATION', bgc=(.9, .7, .3), c=self.loadAnimData)
    
        cmds.setParent('..')
        
        #=====Copy Tab
        self.UIElements["copyTab"] = cmds.rowColumnLayout(nr=7, ral=[(1,'left')],w=100)
        cmds.text(l="First select your source objects and then your destination objects.\nClick 'Copy Animation' to complete.")
        self.UIElements["copySourceList"] = cmds.textScrollList(w=100, h=55)
        self.UIElements["copySource"] = cmds.button(l='Set Source Objects', bgc=(.5,.4,.7), c=self.getSourceObjs)
        cmds.separator()
        self.UIElements["copyDestList"] = cmds.textScrollList(w=100, h=55)
        self.UIElements["copyDesitnation"] = cmds.button(l='Destination Objects',  bgc=(.7,.4,.5), c=self.getDestObjs)
        cmds.separator()
        self.UIElements["copyAnimationBtn"] = cmds.button(l='COPY ANIMATION', bgc=(.9, .7, .3), c=self.copySourceDest)
        
        cmds.setParent('..')
        
        #=====Capture Tab
        self.UIElements["captureTab"] = cmds.rowColumnLayout(nr=3)
        cmds.text("capture realtime animation")
        cmds.button(l="Rec", bgc=(.1,.7,.2))
        cmds.button(l="Stop",bgc=(.7,.2,.1))
        
        cmds.setParent
        
        #rename all the tabs
        cmds.tabLayout(self.UIElements["mainTabs"], e=True, tabLabel=[(self.UIElements["saveTab"],'Save/Load'), (self.UIElements["copyTab"], 'Copy'), (self.UIElements["captureTab"], 'Capture')])
        
        cmds.showWindow(self.UIElements["mainWindow"])
        
    def saveAnimData(self, *args):
        """Select the dir for saving animation data and save the file in XML format"""
        
        sel=[]
        if self.findSelected:
            #sel all objs with animation in scene
            print("sel all")
            sel = cmds.ls(sl=True)
        elif self.findAll:
            #use only Selected objects
            #make sure object has animation on it, no need to use it if not
            for s in sel:
                if cmds.findType(s, type='animCurve'):
                    sel.append(s)
                    print(s +" HAS ANIMATION")
                else:
                    print(s + "DOESN'T HAVE ANIMATION")
        
        #print(len(sel))
        #do some checking to see if items have been selected and they have animation on them to prevent errors
        if len(sel) == 0:
            cmds.confirmDialog(t="No items selected", m="You need to select something dummy!")
            return
        
        for s in sel:
            animCurveNodes = cmds.keyframe(s, q=True, name=True)
            #make sure the object have animation on them
            if animCurveNodes == None:
                cmds.confirmDialog(t="No Animation Data", m="Some or all of your selection does not have animation on it. Aborting.")
                return
        
        #get save location from user
        self.saveDir = cmds.fileDialog2(ds=2, dir=self.rootDir, ff='Zef Anim Files (*.zef)', cap='Save Animation File', fm=0)
        
        #check and see if user selected cancel or not
        if self.saveDir == None:
            return
        
        animNodes = []
        keyTimes = []
        
        startFrame = cmds.textFieldGrp(self.UIElements["saveStartFrame"], q=True, tx=True)
        endFrame = cmds.textFieldGrp(self.UIElements["saveEndFrame"], q=True, tx=True)
        
        
        #build xml to store animation data into
        xFile = xd.Document()
        rootNode = xFile.createElement('AnimDataRoot')
        xFile.appendChild(rootNode)
        
        
        #add a progressBar
        progress = 0
        progBar = cmds.progressWindow(title = 'Saving Animation Data', progress = progress, status= 'Loading: 0%', isInterruptable=True)
            
        for s in sel:
            print("\nCurrent Node: "+ s)
            #Check if the dialog has been canceled
            if cmds.progressWindow(progBar, query=True, isCancelled=True ) :
                break
            
            # Check if end condition has been reached
            if cmds.progressWindow(progBar, query=True, progress=True ) >= 100 :
                break
            
            progress += 5
            
            cmds.progressWindow(progBar, e=True, progress=progress, status=('Loading'+ str(progress)+'%'))
            
            newObj = xFile.createElement('obj')
            newObj.setAttribute('name', s)
            rootNode.appendChild(newObj)
            #need curves for each node
            animNodes = cmds.keyframe(s, q=True, name=True)
            
            for ac in animNodes:
                print("\nAnimation Curve: "+ ac)
                newCurve = xFile.createElement('keyCurve')
                newCurve.setAttribute('name', ac)
                newObj.appendChild(newCurve)
                #get list of all keyTimes
                keyTimes = cmds.keyframe(ac, q=True, timeChange=True)  #this is a list of all the key times
                print("\n"+ str(len(keyTimes)) + " keys: ")
                for kt in keyTimes:
                    #get value of current key
                    #ct = cmds.currentTime(kt, update=False, e=True)
                    kValue = cmds.keyframe(ac, q=True, time=(kt,kt), valueChange=True)
                    print(" [ "+ str(kt) + ","+ str(kValue[0])+"]")
                    
                    #get tangents
                    keyInT = cmds.keyTangent(ac, q=True, inTangentType=True, time=(0,kt))
                    keyOutT = cmds.keyTangent(ac, q=True, outTangentType=True, time=(0,kt))
                    keyInA = cmds.keyTangent(ac, q=True, inAngle=True, time=(0,kt))
                    keyOutA = cmds.keyTangent(ac, q=True, outAngle=True, time=(0,kt))
                    print("("+str(keyInT[0])+" angle="+str(keyInA[0])+","+str(keyOutT[0])+" angle="+str(keyOutA[0])+")")
                    
                    keyFrm = xFile.createElement('key')
                    keyFrm.setAttribute('frame', str(kt))
                    keyFrm.setAttribute('value', str(kValue[0]))
                    keyFrm.setAttribute('inTangent', str(keyInT[0]))
                    keyFrm.setAttribute('outTangent', str(keyOutT[0]))
                    keyFrm.setAttribute('inAngle', str(keyInA[0]))
                    keyFrm.setAttribute('outAngle', str(keyOutA[0]))
                    newCurve.appendChild(keyFrm)
                    
                    #see if key is normal or breakdown
                    #isNorm = cmds.keyframe(ac, q=True, breakdown=True, time=(0,kt))
                    #print("\nKey is: "+str(isNorm[0]))
                    
        print(xFile.toprettyxml())
        #cmds.progressWindow(progBar,endProgress=1)
        #save the file
        #print(self.saveDir)
        f = open(self.saveDir[0], 'w')
        f.write(xFile.toprettyxml())
        f.close()
                 


    def loadAnimData(self, *args):
        """Load the anim data from an XML file"""
        self.loadDir = cmds.fileDialog2(ds=4, dir=self.rootDir, ff='Zef Anim Files (*.zef)', cap='Load Animation File', fm=0)
        
        #check and see if user selected cancel or not
        if self.loadDir == None:
            return
        
        xFile = xd.parse(self.loadDir[0])  #parse XML
        rootNode = xFile.firstChild     #set root node
        
        objects = rootNode.getElementsByTagName('obj')  #get the objects
        for o in objects:
            obj = o.attributes['name'].value   #loop through object and get names
            print("\n"+obj)
            curves = o.getElementsByTagName('keyCurve')   #get the curves
            for c in curves:
                animCurve = c.attributes['name'].value    #loop through curves and get names
                curveVal = str(animCurve).rpartition('_')  #grab just the attr name from the curve
                print("CURVE-VALUE: " + curveVal[2])
                print("curve: "+animCurve)
                keys = c.getElementsByTagName('key')     #get the keys
                for k in keys:
                    frame = k.attributes['frame'].value    #loop through the keys and grab all the attrs
                    value = k.attributes['value'].value
                    inTan = k.attributes['inTangent'].value
                    outTan = k.attributes['outTangent'].value
                    #inAng = k.attributes['inAngle'].value
                    #outAng = k.attributes['outAngle'].value
                    print("frame: "+frame)
                    print("value: "+value)
                    
                    cmds.setKeyframe(obj, itt=inTan, ott=outTan, at=curveVal[2], t=float(frame), v=float(value))   #set keys on the objects and set all the values
                    #cmds.keyframe(obj+'.'+ curveVal[1], edit=True, index=(0,0), timeChange=(frame + 'sec'), valueChange=float(value))
            
            
        #loop through file and get objects and set values
        
    def setRange(self, *args):
        """Enable or Disable the range values for saving"""
        if cmds.textFieldGrp(self.UIElements["saveStartFrame"], q=True, en=True):
            cmds.textFieldGrp(self.UIElements["saveStartFrame"], e=True, en=False)
            cmds.textFieldGrp(self.UIElements["saveEndFrame"], e=True, en=False)
        else:
            cmds.textFieldGrp(self.UIElements["saveStartFrame"], e=True, en=True)
            cmds.textFieldGrp(self.UIElements["saveEndFrame"], e=True, en=True)
        
        self.useRange = cmds.textFieldGrp(self.UIElements["saveStartFrame"], q=True, en=True)
        
    def setType(self, *args):
        """Enable or disable using the type flags"""
        if cmds.optionMenu(self.UIElements["saveTypeMenu"], q=True, en=True):
            cmds.optionMenu(self.UIElements["saveTypeMenu"], e=True, en=False)
        else:
            cmds.optionMenu(self.UIElements["saveTypeMenu"], e=True, en=True)
            
        self.useType = cmds.checkBox(self.UIElements["saveTypeCB"], q=True, en=True)
        
    def storeType(self, *args):
        self.saveByType = cmds.optionMenu(self.UIElements["saveTypeMenu"], q=True, v=True)
       
    def setSaveAll(self, *args):
        self.findSelected = cmds.checkBoxGrp(self.UIElements["saveWhatCB"], q=True, v1=True)
        cmds.checkBoxGrp(self.UIElements["saveWhatCB"], e=True, v2=False)
        
        
    def setSaveSel(self, *args):
        self.findAll = cmds.checkBoxGrp(self.UIElements["saveWhatCB"], q=True, v2=True)
        cmds.checkBoxGrp(self.UIElements["saveWhatCB"], e=True, v1=False)     
        
    def getSourceObjs(self, *args):
        """Get the selected objects and store in a list"""
        
        self.sourceObjects = cmds.ls(sl=True, tr=True)
        
        print("get Source")
    
    def getDestObjs(self, *args):
        """Get the selected objects and store in a list"""
        
        self.destObjects = cmds.ls(sl=True, tr=True)
        print("get Dest")
        
    def copySourceDest(self, *args):
        """Take the source  objects and loop through to get info and copy to dest objects"""
        #how do I match the objects? order of selection?
        
        if len(self.sourceObjects) < 1:
            cmds.confirmDialog(m="You don't have any source objects selected. Please select some first")
            return
        
        if len(self.destObjects) < 1:
            cmds.confirmDialog(m="You don't have any destination objects selected. Please select some first")
            return
        
        print(self.sourceObjects)
        print(self.destObjects)
        

